import os
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.preprocessing import StandardScaler, LabelEncoder
from sklearn.model_selection import train_test_split
from xgboost import XGBClassifier
from sklearn.metrics import classification_report, confusion_matrix
import shap  # 导入 SHAP 库
# 正确的导入方式
from sklearn.preprocessing import MultiLabelBinarizer
from collections import Counter

# 创建输出目录
os.makedirs('output', exist_ok=True)

# 1. 加载数据
stroke = pd.read_csv('../BuildingPrediction/output/stroke_processed.csv')
heart = pd.read_csv('../BuildingPrediction/output/heart_processed.csv')
cirrhosis = pd.read_csv('../BuildingPrediction/output/cirrhosis_processed.csv')

# 2. 统一列名（将 age、gender 等统一为 Age、Sex）
for df in [stroke, heart, cirrhosis]:
    if 'age' in df.columns:
        df.rename(columns={'age': 'Age'}, inplace=True)
    if 'sex' in df.columns:
        df.rename(columns={'sex': 'Sex'}, inplace=True)
    if 'gender' in df.columns:
        df.rename(columns={'gender': 'Sex'}, inplace=True)

# 检查每个表的列名
print("stroke columns:", stroke.columns.tolist())
print("heart columns:", heart.columns.tolist())
print("cirrhosis columns:", cirrhosis.columns.tolist())

# 取三表共有特征（去掉id和标签列）
exclude_cols = {'id', 'stroke', 'HeartDisease', 'Stage', 'label_stroke', 'label_heart', 'label_cirrhosis'}
common_features = list((set(stroke.columns) & set(heart.columns) & set(cirrhosis.columns)) - exclude_cols)
print("最终用于建模的特征：", common_features)

# 构造统一的数据框（不强制要求有id列）
df_stroke = stroke[common_features].copy()
df_heart = heart[common_features].copy()
df_cirrhosis = cirrhosis[common_features].copy()

# 添加标签
df_stroke['label_stroke'] = stroke['stroke']
df_heart['label_heart'] = heart['HeartDisease']
df_cirrhosis['label_cirrhosis'] = (cirrhosis['Stage'] >= 3).astype(int)

# 合并数据
df_combined = pd.concat([df_stroke, df_heart, df_cirrhosis], ignore_index=True)
df_combined.reset_index(drop=True, inplace=True)

# 多标签目标变量 Z
df_combined['multi_label'] = df_combined.apply(lambda row: [row['label_stroke'], row['label_heart'], row['label_cirrhosis']], axis=1)
mlb = MultiLabelBinarizer()
Y = mlb.fit_transform(df_combined['multi_label'])

# 检查多标签组合分布
print("多标签组合分布：", Counter([tuple(row) for row in Y]))

# 只保留出现次数大于1的组合
combo_counts = Counter([tuple(row) for row in Y])
mask = np.array([combo_counts[tuple(row)] > 1 for row in Y])
X = X[mask]
Y = Y[mask]

# 再分层抽样
X_train, X_test, Y_train, Y_test = train_test_split(X, Y, test_size=0.2, random_state=42, stratify=Y)

# 4. 特征工程
X = df_combined[common_features]
X = pd.get_dummies(X)  # 编码分类变量
scaler = StandardScaler()
X_scaled = scaler.fit_transform(X)

# 5. 模型训练
X_train, X_test, Y_train, Y_test = train_test_split(X_scaled, Y, test_size=0.2, random_state=42, stratify=Y)

model = XGBClassifier(random_state=42, eval_metric='logloss')
multi_output_model = MultiLabelBinarizer() # This line was incorrect, should be MultiOutputClassifier
multi_output_model.fit(X_train, Y_train)
Y_pred = multi_output_model.predict(X_test)

# 6. 模型评估
print("多标签分类报告：")
print(classification_report(Y_test, Y_pred, target_names=['Stroke', 'Heart Disease', 'Cirrhosis']))

# 7. SHAP 解释性分析
explainer = shap.TreeExplainer(multi_output_model.estimators_[0])
shap_values = explainer.shap_values(X_test)

# 全局解释
plt.figure(figsize=(10, 8))
shap.summary_plot(shap_values[1], X_test, feature_names=X.columns)
plt.savefig('output/shap_summary.png')
plt.close()

# 局部解释（取前5个样本为例）
for i in range(5):
    plt.figure(figsize=(10, 8))
    shap.force_plot(explainer.expected_value[1], shap_values[1][i], X_test[i], feature_names=X.columns)
    plt.savefig(f'output/shap_force_plot_{i}.png')
    plt.close()

# 8. 特征重要性可视化（以心脏病为例）
feat_imp = pd.Series(multi_output_model.estimators_[1].feature_importances_, index=X.columns)
feat_imp.nlargest(10).plot(kind='barh')
plt.title('心脏病模型特征重要性')
plt.savefig('output/heart_feature_importance.png')
plt.close()

# 9. 共病数量分布
comorbidity_count = Y.sum(axis=1)
sns.countplot(x=comorbidity_count)
plt.title('共病数量分布')
plt.xlabel('共病种数')
plt.ylabel('人数')
plt.savefig('output/comorbidity_count.png')
plt.close()

print("\n✅ 所有分析已完成，图像已保存至 output/ 文件夹")

# 以XGBoost为例
from xgboost import XGBClassifier

# 训练心脏病模型
model_heart = XGBClassifier().fit(X_train, y_heart)
# 训练中风模型
model_stroke = XGBClassifier().fit(X_train, y_stroke)
# 训练肝硬化模型
model_cirrhosis = XGBClassifier().fit(X_train, y_cirrhosis)

class ComorbidityRiskAssessor:
    def __init__(self, models, columns_sets):
        self.models = models
        self.columns_sets = columns_sets

    def _prepare_data(self, patient_dict, disease_name):
        df = pd.DataFrame([patient_dict])
        df_encoded = pd.get_dummies(df)
        training_cols = self.columns_sets[disease_name]
        df_aligned = df_encoded.reindex(columns=training_cols, fill_value=0)
        return df_aligned

    def predict(self, patient_data):
        p_heart = self.models['heart'].predict_proba(self._prepare_data(patient_data, 'heart'))[0, 1]
        p_stroke = self.models['stroke'].predict_proba(self._prepare_data(patient_data, 'stroke'))[0, 1]
        p_cirrhosis = self.models['cirrhosis'].predict_proba(self._prepare_data(patient_data, 'cirrhosis'))[0, 1]
        results = {
            'P(心脏病)': p_heart,
            'P(中风)': p_stroke,
            'P(肝硬化)': p_cirrhosis,
            'P(心脏病 ∩ 中风)': p_heart * p_stroke,
            'P(心脏病 ∩ 肝硬化)': p_heart * p_cirrhosis,
            'P(中风 ∩ 肝硬化)': p_stroke * p_cirrhosis,
            'P(心脏病 ∩ 中风 ∩ 肝硬化)': p_heart * p_stroke * p_cirrhosis
        }
        return results

assessor = ComorbidityRiskAssessor(
    models={'heart': model_heart, 'stroke': model_stroke, 'cirrhosis': model_cirrhosis},
    columns_sets={'heart': cols_heart, 'stroke': cols_stroke, 'cirrhosis': cols_cirrhosis}
)

patient = {
    'Age': 60, 'Sex': 'M', 'bmi': 28, 'Cholesterol': 220, 'avg_glucose_level': 120, 'smoking_status': 'formerly smoked',
    # ... 其他特征 ...
}
risk = assessor.predict(patient)
for disease, prob in risk.items():
    print(f"{disease}: {prob:.2%}")